############################################################
#
# SDSSTools.py - Contains various functions and classes that
#               are relevant to the SDSS
#
############################################################

__author__ = 'Adrian Price-Whelan <adrn@nyu.edu>'

# Standard library dependencies
import sys, os, re

# Third-party imports
import pyfits
import numpy as np
from scipy.ndimage import gaussian_filter1d

# Define all:
#__all__ = ['']

############################################################

def spPlate2Wavelengths(spPlatePath):
    """ Calculate the wavelenths in angstroms from the way the data is stored using
    lambda = 10**(c0 + c1*x) where x is the data as it is stored
    
    @type   spPlateHDUList: HDUList
    @param  spPlateHDUList: The HDU List from an spPlate file from which the wavelengths are calculated
    @rtype: list
    @returns:   A list of wavelength values in units angstroms for this plate's spectra reduction"""
    spPlateHDUList = pyfits.open(spPlatePath)
    npix = spPlateHDUList[0].header['NAXIS1']
    c0 = spPlateHDUList[0].header['COEFF0']
    c1 = spPlateHDUList[0].header['COEFF1']
    spPlateHDUList.close()
    return loglam2wavelengths(npix, c0, c1)

def loglam2wavelengths(npix, c0, c1):
    """ Calculate the wavelengths in angstroms from the way the data is stored using
    lambda = 10**(c0 + c1*x) where x is the data as it is stored
    
    @type   npix: int
    @param  npix: The length of the wavelength grid
    @type   c0: number
    @param  c0: The 1st coefficient using log-lam to convert to wavelengths
    @type   c1: number
    @param  c1: The 2nd coefficient using log-lam to convert to wavelengths
    @rtype: list
    @returns:   A list of wavelength values in units angstroms for this plate's spectra reduction"""
    
    return [10**(c0 + c1*val) for val in range(npix)]

def spPlate2Dict(spfile):
    """ This code was adapted from Ben Weaver's port of readspec to Python. Readspec scans the
    filesystem with every call, so in order to prevent this, I pulled out the necessary code to
    accept a path to a file.
    
    @type   spfile: string
    @param  spfile: The full path to the spPlate file to read in
    @rtype: dict
    @returns:   A dictionary that converts the HDU's from the spPlate file into keys, and the subsequent data as values
    """
    
    spplate = pyfits.open(spfile)
    npix = spplate[0].header['NAXIS1']
    c0 = spplate[0].header['COEFF0']
    c1 = spplate[0].header['COEFF1']
    
    thisfiber = np.array(spplate[5].data.field("FIBERID"))
    
    coeff0 = np.zeros(thisfiber.size,dtype='d') + c0
    coeff1 = np.zeros(thisfiber.size,dtype='d') + c1
    loglam0 = c0 + c1*np.arange(npix,dtype='d')
    loglam = np.resize(loglam0,(thisfiber.size,npix))
    
    spplate_data = dict()
    hdunames = ('flux','invvar','andmask','ormask','disp','plugmap','sky','loglam',)
    
    spfile_parse = os.path.basename(spfile).split("-")
    plate = int(spfile_parse[1])
    mjd = int(spfile_parse[2].split(".")[0])
    
    platevec = np.zeros(len(thisfiber),dtype='i4') + plate
    mjdvec = np.array([mjd],dtype='i4')
    pmjdindex = np.where(thisfiber > 0, thisfiber-1, 0) # ????

    #
    # Read the data images
    #
    for k in range(len(hdunames)):
        try:
            tmp = spplate[k].data[np.where(thisfiber > 0, thisfiber-1, 0),:]
        except IndexError:
            tmp = loglam
        if hdunames[k] not in spplate_data:
            if k == 0:
                allpmjdindex = pmjdindex
                allcoeff0 = coeff0
                allcoeff1 = coeff1
            #
            # Put the data into the return structure
            #
            if hdunames[k] == 'plugmap':
                spplate_data['plugmap'] = dict()
                for c in spplate[5].columns.names:
                    spplate_data['plugmap'][c] = tmp.field(c)
            else:
                spplate_data[hdunames[k]] = tmp
        else:
            #
            # Append data
            #
            if k == 0:
                allpmjdindex = np.concatenate((allpmjdindex,pmjdindex))
                ps = 0
                allcoeff0 = np.concatenate((allcoeff0,coeff0))
                allcoeff1 = np.concatenate((allcoeff1,coeff1))
            if hdunames[k] == 'plugmap':
                for c in spplate[5].columns.names:
                    spplate_data['plugmap'][c] = np.concatenate(
                        (spplate_data['plugmap'][c],tmp.field(c)))
            else:
                spplate_data[hdunames[k]] = spec_append(spplate_data[hdunames[k]],tmp,pixshift=ps)
    spplate.close()
    
    #
    # Reorder the data.  At this point allpmjdindex is an index for which
    # fiber[allpmjdindex] == spplate['plugmap']['FIBERID'], so we have to
    # reverse this mapping.
    #
    j = allpmjdindex.argsort()
    for k in spplate_data:
        if isinstance(spplate_data[k],dict):
            for c in spplate_data[k]:
                if spplate_data[k][c].ndim == 2:
                    spplate_data[k][c] = spplate_data[k][c][j,:]
                else:
                    spplate_data[k][c] = spplate_data[k][c][j]
        else:
            spplate_data[k] = spplate_data[k][j,:]
    allcoeff0 = allcoeff0[j]
    allcoeff1 = allcoeff1[j]
    #
    # If necessary, recompute the wavelengths
    #
    nfibers,npixmax = spplate_data['flux'].shape

    return spplate_data
    
def textList2List(text):
    """ Parses a string input of numbers to get a python list of all the numbers
    
    Accepted formats:
        Range : 15-20
        List : 1,2,3,4,5
        Combination : 1,2,3,15-20
    Output:
        Range : [15,16,17,18,19,20]
        List : [1,2,3,4,5]
        Combination : [1,2,3,15,16,17,18,19,20]
    
    @type   text: string
    @param  text: Input text that represents a list of numbers
    @rtype: list
    @returns:   list of integers
    """
    
    text = str(text)
    
    textList = text.replace(" ","")
    textList = textList.split(",")
    
    returnList = []
    for tmp in textList:
        minmax = tmp.split("-")
        if len(minmax) == 2:
            num_range = range(int(minmax[0]),int(minmax[1])+1)
            returnList += num_range
        else:
            returnList += [int(minmax[0])]
    returnList.sort()
    
    return list(np.unique(returnList))

def smooth_spectrum(wave, flux, sigma):
    """ Takes a wavelength array and flux array and smooths/downsamples them
    @type   wave: list
    @param  wave: Input a list of wavelength values for the flux
    @type   flux: list
    @param  flux: A spectral flux
    @type   sigma: float
    @param  sigma: Sigma in a Gaussian smooth of the input arrays
    @rtype: list
    @returns:   list of arrays: wavelength, flux
    """
    degree = int(8*sigma + 1)
    return wave[::int(degree/2)], list(gaussian_filter1d(flux, sigma)[::int(degree/2)])
    
def parsePMF(pmf, delimiters=[",","|"," ","-"]):
    """ This function is meant to parse a string Plate, MJD, Fiber
        and return the 3 numbers as integers
    
    Parameters
    ----------
    pmf : string
        A string containing Plate ID, MJD, Fiber ID delimited by one 
        of the specified delimiters.
    delimiters : list, string
        Delimiters to try when splitting the string.
    
    Returns
    -------
    split : tuple
        Returns a tuple of integers (Plate, MJD, Fiber)
        
    """
    for delim in delimiters:
        split = pmf.split(",")
        if len(split) == 3: return map(int, split)

def parsePMFs(pmf):
    """ This function is meant to parse a string Plate, MJD, and list
        of Fibers
    
    Parameters
    ----------
    pmf : string
        A string containing Plate ID, MJD, Fiber IDs in the form
        plate,mjd,[fibers] where [fibers] can be of the form:
            [1-20,22,25]
    delimiters : list, string
        Delimiters to try when splitting the string.
    
    Returns
    -------
    split : tuple
        Returns a tuple of integers (Plate, MJD, Fiber)
        
    """
    pattr = re.compile("([0-9]{4})[,|]{0,1}\s*([0-9]{5})[,|]{0,1}\s*\[*([0-9-\,]*)\]*",re.UNICODE)
    match = pattr.search(pmf)
    plate, mjd, fibers = match.groups()
    
    if fibers != "":
        parsedFibers = textList2List(fibers)
    else:
        parsedFibers = []
    
    return plate, mjd, parsedFibers
